package com.zjj.lbw.ai;

import com.google.common.collect.Lists;
import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.zhipu.ZhipuAiChatModel;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.time.LocalDateTime;
import java.util.Collections;

public class ToolsTest {
    @Tool("获取当前日期")
    public static String dateUtil(String onUse) {
        return LocalDateTime.now().toString();
    }

    public static void main(String[] args) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {

        ChatLanguageModel model = ZhipuAiChatModel.builder()
                .apiKey("9a6dfc4258104e28ad0c72db2ddda3e4.f6VL6AWRQxORuJpi")
                .build();


        ToolSpecification toolSpecification = ToolSpecifications.toolSpecificationFrom(
                ToolsTest.class.getMethod("dateUtil", new Class[]{String.class}));

        UserMessage userMessage = UserMessage.from("今天是几月几号？");

//      所以，一个ToolSpecification对象就代表一个工具，当用户把要问题UserMessage和工具ToolSpecification一起传递给大模型，
//      大模型就知道要结合工具描述来解决用户的问题，此时大模型响应的AiMessage不再是一串文本，而是一个工具执行请求
        Response<AiMessage> response = model.generate(Collections.singletonList(userMessage), toolSpecification);

        System.out.println(response.content());

        // 我们在得到了ToolExecutionRequest后，就需要取执行对应的工具方法了
        AiMessage aiMessage = response.content();
        if (aiMessage.hasToolExecutionRequests()) {
            for (ToolExecutionRequest toolExecutionRequest : aiMessage.toolExecutionRequests()) {
                String methodName = toolExecutionRequest.name();
                Method method = ToolsTest.class.getMethod(methodName, new Class[]{String.class});

                // result就是当前时间
                String result = (String) method.invoke(null, "123");
                System.out.println(result);

                // 将工具执行结果 返回给大模型
                ToolExecutionResultMessage executionResultMessage = ToolExecutionResultMessage.from(
                        toolExecutionRequest.id(), toolExecutionRequest.name(), result
                );
                AiMessage message = model.generate(Lists.newArrayList(userMessage, aiMessage, executionResultMessage)).content();
                System.out.println(message.text());
            }
        }

    }
}
